#AES加密
import box
import key
import GFarith
import exsystem
import array

S = box.getS()
IS = box.getIS()

'''kk = ['0f','15','71','c9','47','d9','e8','59','0c','b7','ad','d6','af','7f','67','98']

akey = key.generate(kk)'''

matrix = {}
matrix[0] = ['02','03','01','01']
matrix[1] = ['01','02','03','01']
matrix[2] = ['01','01','02','03']
matrix[3] = ['03','01','01','02']

imatrix = {}
imatrix[0] = ['0e','0b','0d','09']
imatrix[1] = ['09','0E','0B','0D']
imatrix[2] = ['0D','09','0E','0B']
imatrix[3] = ['0B','0D','09','0E']



#列混窑变换
def Maxchange(new, matrix):
    tmp = {}
    result = []
    for i in range(4):
        tmp[i] = new[i*4:i*4+4]
    for i in range(4):
        for j in range(4):
            ele = '00'
            for k in range(4):
                ele = GFarith.hexplus(ele, GFarith.hexmulti(matrix[i][k],tmp[k][j]))
            result.append(ele)
    return result

#逆列混淆变换
def Imaxchange(new, imatrix):
    tmp = {}
    result = []
    for i in range(4):
        tmp[i] = new[i*4:i*4+4]
    for i in range(4):
        for j in range(4):
            ele = '00'
            for k in range(4):
                ele = GFarith.hexplus(ele, GFarith.hexmulti(imatrix[i][k],tmp[k][j]))
            result.append(ele)
    return result

                
#明文以列表形式输入，16个字节一组
def turn(plain,key):
    plain = reverse(plain)
    plain = Schange(plain)
    plain = rowmove(plain)
    pl = Maxchange(plain, matrix)
    pl = reverse(pl)
    result = [GFarith.hexplus(pl[i],key[i]) for i in range(16)]
    return result

def iturn(ciper, key):
    ciper = reverse(ciper)
    ciper = irowmove(ciper)
    ciper = ISchange(ciper)
    key = reverse(key)
    ci = [GFarith.hexplus(ciper[i],key[i]) for i in range(16)]
    ci = Imaxchange(ci,imatrix)
    ci = reverse(ci)
    return ci

def reverse(n):
    new = []
    for i in range(4):
        for j in range(4):
            new.append(n[i + j*4])
    return new

def Schange(n):
    for i in range(len(n)):
        a = int(n[i][0],16)
        b = int(n[i][1],16)
        n[i] = S[a][b]
    return n

def ISchange(n):
    for i in range(len(n)):
        a = int(n[i][0],16)
        b = int(n[i][1],16)
        n[i] = IS[a][b]
    return n

def rowmove(new):
    new[4:8] = new[5:8] + new[4:5]
    new[8:12] = new[10:12] + new[8:10]
    new[12:16] = new[15:16] + new[12:15]
    return new

def irowmove(new):
    new[4:8] = new[7:8] + new[4:7]
    new[8:12] = new[10:12] + new[8:10]
    new[12:16] = new[13:16] + new[12:13]
    return new

def wholeturn(plain,akey):
    #plain = [exsystem.detohex(plain[i]) for i in range(16)]
    key = akey[0] + akey[1] + akey[2] + akey[3]
    plain = [GFarith.hexplus(plain[i],key[i]) for i in range(16)]
    for i in range(1,10):
        key = akey[i*4] + akey[i*4 + 1] + akey[i*4 + 2] + akey[i*4 + 3]
        plain = turn(plain,key)
    plain = reverse(plain)
    plain = Schange(plain)
    plain = rowmove(plain)
    plain = reverse(plain)
    key = akey[40] + akey[41] + akey[42] + akey[43]
    plain = [GFarith.hexplus(plain[i],key[i]) for i in range(16)]
    return plain

def iwholeturn(ciper,akey):
    key = akey[40] + akey[41] + akey[42] + akey[43]
    ciper = [GFarith.hexplus(ciper[i],key[i]) for i in range(16)]
    for i in range(1,10):
        key = akey[40 - i*4] + akey[41 - i*4] + akey[42 - i*4] + akey[43 - i*4]
        ciper = iturn(ciper,key)
    key = akey[0] + akey[1] + akey[2] + akey[3]
    ciper = reverse(ciper)
    ciper = irowmove(ciper)
    ciper = ISchange(ciper)
    ciper = reverse(ciper)
    ciper = [GFarith.hexplus(ciper[i],key[i]) for i in range(16)]
    return ciper


'''def testturn(plain,key):
    new = [int(i,16) for i in plain]
    return turn(new,key)'''

def encrypt():
    plainame = input("Please enter the name of plaintext: ")
    cipername = input("Please enter the name of ciper text: ")
    keyname = input("Please enter the name of key: ")
    akey = createkey(keyname)
    with open(plainame,'rb') as f,open(cipername,'wb') as cp:
        pl = f.read(16)
        while pl:
            pl = list(pl)
            if len(pl) == 16:
                pl = [exsystem.detohex(pl[i]) for i in range(16)]
                pl = wholeturn(pl,akey)
            else:
                pl = addcontain(pl)
                pl = [exsystem.detohex(pl[i]) for i in range(16)]
                pl = wholeturn(pl,akey)
            pl = [int(ele,16) for ele in pl]
            pl = array.array('B',pl)
            cp.write(pl)
            pl = f.read(16)

def decode():
    cipername = input("Please input the name of ciper text: ")
    plainame = input("Please input the name of plain text: ")
    keyname = input("Please input the name of key: ")
    akey = createkey(keyname)
    with open(cipername,'rb') as f,open(plainame,'wb') as cp:
        pl = f.read(16)
        while pl:
            pl = list(pl)
            if len(pl) == 16:
                pl = [exsystem.detohex(pl[i]) for i in range(16)]
                pl = iwholeturn(pl,akey)
            else:
                pl = addcontain(pl)
                pl = [exsystem.detohex(pl[i]) for i in range(16)]
                pl = iwholeturn(pl,akey)

            pl = [int(ele,16) for ele in pl]
            pl = array.array('B',pl)
            cp.write(pl)
            pl = f.read(16)
            


def addcontain(pl):
    e = 16 - len(pl)
    while len(pl) < 16:
        pl.append(e)
    return pl


def createkey(keyname):
    k = []
    with open(keyname,'r') as f:
        ch = f.read(2)
        while ch:
            k.append(ch)
            ch = f.read(2)
    akey = key.generate(k)
    return akey
    

